import glob
import importlib

from timm.layers import set_layer_config
from timm.models._hub import load_model_config_from_hf
from utils.registry import Registry
MODEL = Registry('Model')

def get_model(cfg_model):
    # 提取模型名称
    model_name = cfg_model.get('name', '')  

    # 适配新的 dict 参数格式，确保 kwargs 是字典
    kwargs = cfg_model.get('kwargs', {})  # 使用 .get() 方法获取，避免 KeyError
    if not isinstance(kwargs, dict):
        raise ValueError("Expected 'kwargs' to be a dictionary, but got type {}".format(type(kwargs)))

    # 获取模型函数
    model_fn = MODEL.get_module(model_name)
    
    # 处理 timm_ 开头的模型
    if model_name.startswith('timm_'):
        model_name_hf = kwargs.pop('hf', None)
        pretrained_cfg = None
        
        if model_name_hf:
            pretrained_cfg, model_name_hf = load_model_config_from_hf(model_name_hf)
            pretrained_cfg['url'] = ''
        
        with set_layer_config(scriptable=None, exportable=None, no_jit=None):
            model = model_fn( pretrained_cfg=pretrained_cfg, **kwargs)

    else:
        model = model_fn(**kwargs)

    return model

files = glob.glob('models/[!_]*.py')
for file in files:
    model_lib = importlib.import_module(file.split('.')[0].replace('/', '.'))
